import tensorflow as tf
import os

# 设置GPU内存按需增长
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

basePath = "/data/wengjj/catdog"
basePath = "/Users/weng/PycharmProjects2/catdog"


# 加载数据
def load_data():
    train_dir = r'' + basePath + '/data'
    test_dir = r'' + basePath + '/data'
    # 训练集
    train_cats_path = os.path.join(train_dir, 'cats')
    train_dogs_path = os.path.join(train_dir, 'dogs')
    # 测试集
    test_cats_path = os.path.join(test_dir, 'cats')
    test_dogs_path = os.path.join(test_dir, 'dogs')

    # 定义图像尺寸
    img_height, img_width = 150, 150
    batch_size = 32

    # 数据预处理
    train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
        rescale=1. / 255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True)

    test_datagen = tf.keras.preprocessing.image.ImageDataGenerator(rescale=1. / 255)

    train_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=(img_height, img_width),
        batch_size=batch_size,
        class_mode='binary')

    validation_generator = test_datagen.flow_from_directory(
        test_dir,
        target_size=(img_height, img_width),
        batch_size=batch_size,
        class_mode='binary')

    return train_generator, validation_generator


# 构建模型
def create_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2, 2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(512, activation='relu'),
        tf.keras.layers.Dense(1, activation='sigmoid')
    ])

    model.compile(loss='binary_crossentropy',
                  optimizer=tf.keras.optimizers.Adam(lr=1e-4),
                  metrics=['accuracy'])

    return model


# 训练模型
def train_model(model, train_generator, validation_generator):
    epochs = 300
    steps_per_epoch = train_generator.n // train_generator.batch_size
    validation_steps = validation_generator.n // validation_generator.batch_size

    history = model.fit(
        train_generator,
        steps_per_epoch=steps_per_epoch,
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=validation_steps)

    # 将训练好的模型保存到文件中
    model.save(r'' + basePath + '/cat_dog_classifier300.h5')

    return history


if __name__ == '__main__':
    train_generator, validation_generator = load_data()
    print(train_generator)
    print(validation_generator)
    model = create_model()
    history = train_model(model, train_generator, validation_generator)
